import psycopg2
import re
import time

from simple_naive_bayes import naivebayes,getwords,ident

con = psycopg2.connect(host='dw02.dev01.corp.metaweb.com', database='common',
                       user='postgres')
cur=con.cursor()

category_list=['Category:Ninjas',
               'Category:Pirates',
               'Category:Assassins',
               'Category:Apple Inc. employees',
               'Category:Microsoft employees',
               'Category:Google employees',
               'Category:Free software programmers',
               'Category:Computer programmers',
               'Category:American criminals'
              ]
print "Training classifiers..."

def get_category_members(cur, category):
    records = []
    queue = [category]
    recordsSeen = set()
    while len(queue) > 0 and len(records) < 500:
        currentCategory = queue.pop(0)
        cur.execute("select articles.wpid, articles.name from wikipedia.category_members, wikipedia.articles "+
                    "where category_members.category_name like %s and articles.wpid=category_members.article_wpid", (currentCategory,))
        result = cur.fetchall()
        for wpid, name in result:
            if wpid not in recordsSeen:
                recordsSeen.add(wpid)
                if name.startswith("Category:"): 
                    queue.append(name)
                else: records.append(name)
    return records

# Get the members of every category
name_classes={}
for category in category_list:
    rec = get_category_members(cur, category)
    print str(len(rec)) + " examples for " + category
    for name in rec:
        name_classes.setdefault(name,set()).add(category)

# Test sets:
test_set=["Long John Silver", "Jack Sparrow", "Storm Shadow (G.I. Joe)", "Leonardo (TMNT)", 
          "Bill Gates", "Steve Jobs", "Richard Stallman", "Larry Page", "Guido van Rossum", "Larry Wall", "Jerry Yang", 
          "Hans Reiser", "Theodore Kaczynski", "Kenneth Lay"]

# Remove test sets before training
for name in test_set:
    if name in name_classes: del name_classes[name]

def get_article_text(cur, name):
    cur.execute("select name, text from wikipedia.articles where name=%s", (name,))
    return cur.fetchone()

# Train a classifier for each class
classifiers=[(cat,naivebayes(ident)) for cat in category_list]

for name in name_classes:
    name, text = get_article_text(cur, name)
   
    #print name
    words=getwords(text[0:1024])
    
    for cat,cl in classifiers:
        if cat in name_classes[name]:
            cl.train(words,1)
        else:
            cl.train(words,0)

# Run tests:
for testName in test_set:
    name, text = get_article_text(cur, testName)
    if name.startswith('Category:'): continue
    
    print name
    words=getwords(text[0:1024])
    
    for cat,cl in classifiers:
        py,pn=cl.prob(words,1),cl.prob(words,0)
        print '%s\t%s\t%f' % (cat,cl.classify(words),py/pn if pn>0 else 100)

    print

cur.close()
con.close()

# Try building a string that will fall into multiple categories
test_string = "Toby Segaran lived during the sengoku period in Japan. He spent many years at sea battling big Japanese ships."
print test_string
words=getwords(test_string)
for cat,cl in classifiers:
    py,pn=cl.prob(words,"yes"),cl.prob(words,"no")
    print '%s\t%s\t%f' % (cat,cl.classify(words),py/pn if pn>0 else 100)

